package com.mercycrazy.recruit.filter.xss;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.http.HtmlUtil;
import com.alibaba.fastjson2.JSON;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.StringUtils;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.*;

/**
 * Xss HttpServletRequest Wrapper
 *
 * @author <a href="mailto:1443424326@qq.com">mercycrazy</a>
 * @see HttpServletRequestWrapper
 * @since 2023-01-08
 */
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

    public XssHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
    }

    @Override
    public String getParameter(String name) {
        String value = super.getParameter(name);
        return returnCleanTagValue(value);
    }

    @Override
    public String[] getParameterValues(String name) {
        String[] values = super.getParameterValues(name);
        return returnCleanTagValues(values);
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        Map<String, String[]> parameters = super.getParameterMap();
        Map<String, String[]> newParameters = new LinkedHashMap<>();
        if (null != parameters) {
            parameters.keySet()
                    .forEach(key -> {
                        String[] values = parameters.get(key);
                        String[] returnValues = returnCleanTagValues(values);
                        newParameters.put(key, returnValues);
                    });
        }
        return newParameters;
    }

    @Override
    public String getHeader(String name) {
        String value = super.getHeader(name);
        return returnCleanTagValue(value);
    }

    @Override
    public Enumeration<String> getHeaders(String name) {

        List<String> values = Collections.list(super.getHeaders(name));

        List<String> returnValues = Arrays.asList(
                returnCleanTagValues(values.toArray(new String[0]))
        );

        return Collections.enumeration(returnValues);
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {

        String body = IOUtils.toString(
                super.getInputStream(), StandardCharsets.UTF_8.name());

        Map<String, Object> originMap = BeanUtil.beanToMap(JSON.parse(body));
        Map<String, Object> resultMap = new LinkedHashMap<>();

        // originMap -> resultMap
        originMap.keySet().forEach(key -> {
            Object value = originMap.get(key);

            if (value instanceof String) {
                value = returnCleanTagValue(value.toString());
            }

            resultMap.put(key, value);
        });

        String resultJson = JSON.toJSONString(resultMap);

        ByteArrayInputStream byteArrayInputStream =
                new ByteArrayInputStream(resultJson.getBytes(StandardCharsets.UTF_8));

        return new ServletInputStream() {
            @Override
            public boolean isFinished() {
                return false;
            }

            @Override
            public boolean isReady() {
                return false;
            }

            @Override
            public void setReadListener(ReadListener readListener) {

            }

            @Override
            public int read() throws IOException {
                return byteArrayInputStream.read();
            }
        };
    }

    /**
     * 返回 XSS 转义后的 Value
     *
     * @param value String 类型的 Value
     */
    private String returnCleanTagValue(String value) {
        String returnValue = value;
        if (StringUtils.isNotBlank(returnValue)) {
            returnValue = HtmlUtil.cleanHtmlTag(returnValue);
        }
        return returnValue;
    }

    /**
     * 返回 XSS 转义后的 Value
     *
     * @param values String[] 类型的 Value
     */
    private String[] returnCleanTagValues(String[] values) {

        String[] returnValues = null;
        if (null != values) {
            returnValues = Arrays.stream(values)
                    .map(this::returnCleanTagValue)
                    .toArray(String[]::new);
        }
        return returnValues;
    }
}
